import os

import cv2
from ultralytics import YOLO


def process_image(model, image_path, output_folder, table_output_folder, threshold=0.9):
    try:
        # 读取原始图片
        original_image = cv2.imread(image_path)
        if original_image is None:
            print(f"无法读取图片: {image_path}")
            return
        height, width = original_image.shape[:2]
        # 进行预测
        results = model.predict(source=image_path, conf=threshold)
        table_count = 0
        # 处理预测结果
        for result in results:
            boxes = result.boxes  # 检测框信息
            for i, box in enumerate(boxes):
                class_id = int(box.cls[0])  # 类别 ID
                confidence = float(box.conf[0])  # 置信度
                x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)  # 边界框坐标
                print(
                    f"图片 {os.path.basename(image_path)}: 检测到类别 ID: {class_id}，置信度: {confidence:.2f}，边界框坐标: ({x1}, {y1}, {x2}, {y2})")
                # 扩展坐标并处理越界问题
                new_y1 = max(0, y1 - 20)
                new_y2 = min(height, y2 + 20)
                new_x1 = max(0, x1 - 20)
                new_x2 = min(width, x2 + 20)
                # 截取表格区域
                table_image = original_image[new_y1:new_y2, new_x1:new_x2]
                table_output_path = os.path.join(table_output_folder, f"{table_count}_{os.path.basename(image_path)}")
                try:
                    cv2.imwrite(table_output_path, table_image)
                    table_count += 1
                except Exception as e:
                    print(f"保存表格图片时出错: {table_output_path}, 错误信息: {e}")
            # 构建输出图片的完整路径
            output_path = os.path.join(output_folder, os.path.basename(image_path))
            try:
                # 保存可视化后的图片
                annotated_image = result.plot()
                cv2.imwrite(output_path, annotated_image)
            except Exception as e:
                print(f"保存可视化图片时出错: {output_path}, 错误信息: {e}")
    except Exception as e:
        print(f"处理图片 {image_path} 时出错: {e}")


def main():
    # 加载训练好的模型
    model = YOLO('runs/detect/phase2_finetune/weights/last.pt')
    # 输入图片文件夹路径
    input_folder = '../images'
    # 输出结果保存文件夹路径
    output_folder = '../images/results'
    # 输出表格截取图片的文件夹路径
    table_output_folder = '../images/table_results'
    # 如果输出文件夹不存在，则创建它
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    # 检查表格输出文件夹是否存在，若不存在则创建
    if not os.path.exists(table_output_folder):
        os.makedirs(table_output_folder)
    # 遍历输入文件夹下的所有图片文件
    for filename in os.listdir(input_folder):
        if filename.endswith(('.png', '.jpg', '.jpeg')):
            # 构建图片的完整路径
            image_path = os.path.join(input_folder, filename)
            process_image(model, image_path, output_folder, table_output_folder, threshold=0.8)
    print("所有图片预测完成，结果已保存到指定文件夹。")


if __name__ == "__main__":
    main()
